
%% Using Psychophysics Toolbox Version 3 (PTB-3)
% Toolbox: Dave Brainard, Denis Pelli, Mario Kleiner, et al.
% Experiment: Patrick Connor and Leah Brown
% July/August 2011

delete meta_data.mat
close all; clear all; clc;

%% experiment set up

n_rect=200; % number of colored and gray rectangles displayed on screen
n_phases=1; % number of test phases
n_trials=[20,20,20]; % number of trials for phase 1, 2, etc.
print_results=1; % print experiment trials to command window?
save_results=1; % save experiment results to .txt file?

% reward variables
p_reward_nwr=0; % probability of giving "no reward when reward" (nwr) should be given
p_reward_rwn=0; % probability of giving "reward when no reward" (rwn) should be given
p_r_state=0.5; % probability of a rewarded state shown [don't change]

% don't change these set up variables!!
n_colors=9; % total number of colors availible to choose from
colors_shown=3; % number of colors to be shown to the user at once
Rexp=0; % expected reward
reward=0; % reward actually given (according to p_reward_nwr or p_reward_rwn)


%% set up window
%[window,rect]=Screen('OpenWindow',0,[140, 140, 140]);
Screen('Preference','SkipSyncTests',1);
Screen ('Preference','VisualDebugLevel',3);
[window,rect]=Screen('OpenWindow',1,[140, 140, 140]);

% show arrow cursor -> easiest to see against gray background.
if strcmp(OSName,'Windows') || strcamp(OSName,'OSX'); ShowCursor('Arrow');
else ShowCursor(5); end; clc;

% adjust everything to current computer monitor size
% L,T,B and dL -> landmark measurement to put things on screen.
L=round((rect(3)-rect(1))*0.25);  T=round(0.25*(rect(4)-rect(2))); B=round(0.6*(rect(4)-rect(2)));
dL=round((1/n_rect)*(0.5*(rect(3)-rect(1))));

% font size for small and large text. May need to be changed depending on
% size of monitor
small_text=round((rect(3)-rect(1))/50);
big_text=round((rect(3)-rect(1))/35);

% left and right hoizontal limits for rating number placement [don't
% change]
dL2=(0.4*(rect(3)-rect(1))-10*20)/(10);
rect_box=[0.3*(rect(3)-rect(1))+(0:9)*(dL2+20);0.3*(rect(3)-rect(1))+(1:10)*(dL2)+(0:9)*20;];
Screen(window,'Flip'); Screen('TextSize', window ,small_text);

%% colors to be used on screen.

% c_names: yellow, green, blue, violet, pink, red, brown, white, black
c_names='YGBVPRNWK';
% RGB values for those colors: yellow, green, blue, violet, pink, red, brown, white, black
col_spec=[255,255,0; 0,250,0; 0,191,255; 155,48,255; 255,150,190; 255,0,0; 139,69,0; 255,255,255; 0,0,0]';

% now scramble colors to randomly generate colors of stimuli A,B,C,D,...,X
temp=randperm(n_colors); stimA=1; stimB=2; stimC=3; stimD=4; stimE=5; stimX=n_colors;
c_n=c_names(temp); cols=col_spec(:,temp); cols(:,n_colors+1)=[140, 140, 140]';

str=1; % iterator for recording trials

%% data recording variables
% guesses: matrix holding guesses from post phase prediction period.
% indexed by [phase, color], where color is in order of c_n
% padded by 0s
guesses=zeros(n_phases,n_colors);
% ratings: matrix holding ratings from each trial
% indexed by [trial, phase]
% padded by 0s
ratings=zeros(n_trials, n_phases);

% r_exp_record and r_giv_record: records of expected and given reward
% padded by 2
r_exp_record=zeros(n_trials,n_phases)+2;
r_giv_record=zeros(n_trials,n_phases)+2;


%% GROUP SELECTION
group=2;

%% printing results
if print_results
    disp('               COLOR KEY')
    disp('      Y yellow       G green        B blue');
    disp('      V violet       P pink         R red');
    disp('      N brown        W white        K black');
    disp('================ EXPERIMENT RESULTS =================');
    disp(['Stimulus A: ',c_n(1), '          Stimulus B: ',c_n(2)]);
    disp(['Stimulus C: ',c_n(3), '          Stimulus D: ',c_n(4)]);
    disp(['Stimulus E: ',c_n(5), '          Stimulus F: ',c_n(6)]);
    disp(['Stimulus G: ',c_n(7), '          Stimulus H: ',c_n(8)]);
    disp(['Stimulus X: ',c_n(9)]);
end

%% BEGIN EXPERIMENT
for phase=1:n_phases
    % display phase
    Screen('TextSize', window, big_text);
    DrawFormattedText(window, sprintf('Phase %d',phase),'center','center',[255,255,255]);
    Screen(window,'Flip'); Screen('TextSize', window, small_text);
    pause(2);
    
    % recording data
    
    trials{str}= sprintf('\n\n\t\t\t\tPHASE %d:\n', phase); %#ok<*SAGROW>
    trials{str+1}=sprintf('TRL\tRexp\tRgiv\tRating\tColors');
    if print_results
        fprintf('\n\n\t\tPHASE %d:\n%s\n', phase,trials{str+1});
    end
    str=str+2;
    
    
    % show_r_trials -> selected trial numbers to show rewarded stimuli
    show_r_trials=randperm(n_trials(phase));
    show_r_trials=show_r_trials(1:round(p_r_state*n_trials(phase)));
    % show_AX_trials -> selected trial numbers to show AX stimuli (for
    % phase 1)
    AX_trials=show_r_trials(1:floor(length(show_r_trials(phase))/2));
    
    for trial=1:n_trials(phase)
        
        %% present stimuli ------------------------------------------------
        Rexp=0; reward=rand<p_reward_rwn;
        switch phase
            % for each trial, select three colors to be shown. Check to see
            % if any of the chosen colors are duplicates or an excluded
            % color. If so, reselect colors. If the trial is a rewarded
            % one, change reward to reflect this.
            case 1
                % ??? -> exclude A or B
                colors=unique(randi(n_colors-3,[1, colors_shown])+2);
                while any(colors==5) || length(colors)<3
                    colors=unique(randi(n_colors-3,[1, colors_shown])+2);
                end
                % [A/B]?? -> exclude E
                if any(show_r_trials==trial);
                    if any(AX_trials==trial)
                        colors(1)=1;
                    else colors(1)=2;
                    end
                    colors(2)=n_colors;
                    Rexp=1; reward=rand>p_reward_nwr;
                end
            case 2
                % group 1: A?? -> exclude A,B,C,D,X
                % group 2: E?? -> exclude B,C,D,E,X
                tempcol=unique(randi(3,[1,2])+5);
                while(length(tempcol)<2); tempcol=unique(randi(3,[1,2])+5); end
                colors(2:3)=tempcol;
                if group==1; colors(1)=1;
                else colors(1)=5; end
                % CD? -> exclude A,B,X,E
                if any(show_r_trials==trial);
                    colors(1:2)=[3,4];
                    Rexp=1; reward=rand>p_reward_nwr;
                end
            otherwise
                % B?? -> exclude C, D, X, A, E.
                colors(1)=2;
                tempcol=unique(randi(3,[1,2])+5);
                while length(tempcol)<2; tempcol=unique(randi(3,[1,2])+5); end
                colors(2:3)=tempcol;
        end
        
        % record reward
        r_exp_record(trial,phase)=Rexp;
        r_giv_record(trial,phase)=reward;
        
        % now disperse chosen colors among the stripes
        % don't allow consecutive same colored stripes.
        color=colors(randi(colors_shown,1,2*n_rect));
        sep_color=color(diff(color)~=0);
        clear color;
        color(1:2:n_rect)=sep_color(1:n_rect/2);
        color(2:2:n_rect)=n_colors+1; % 10th color is gray.
        
        %% training - receive and analyze input----------------------------
        % rate_box holds bounding box of rating numbers so they can be
        % clicked on.
        rate_box=zeros(10,4);
        rating=[];
        % display stimuli
        Screen('FillRect',window,cols(:,color),[L+(0:n_rect-1)*dL; T*ones(1,n_rect); L+(1:n_rect)*dL; B*ones(1,n_rect)]);
        [~,~,bounds]=DrawFormattedText(window, 'Rating of how likely this is to receive reward:','center',B+small_text,[255,255,255]);
        for j=1:10; [~,~,rate_box(j,:)]=DrawFormattedText(window,sprintf('%d',j),(rect_box(1,j)+rect_box(2,j))/2,bounds(4)+small_text,[190, 190, 190]); end
        Screen(window,'Flip',[],1); % don't clear background
        
        while isempty(rating) % until a rating is made
            
            % wait for a mouse click
            [x,y,clicks]=GetMouse(window);
            while ~clicks; [x,y,clicks]=GetMouse(window); end
            
            % check if the mouse click is in the bounding box of a rating
            % number, if so find which one.
            
            rating=find(sum(rate_box(:,[2 4])>=y,2)==1 & sum(rate_box(:,[1 3])>=x,2)==1); %#ok<*VUNUS>
            
            %            if y>=rate_box(2,2) && y<=rate_box(2,4)
            %                 for i=1:10;
            %                     if x>=rate_box(i,1) && x<=rate_box(i,3)
            %                         rating=i;
            %                         break;
            %                     end
            %                 end
            %            end
        end %        while isempty(rating) % until a rating is made
        % if there is a rating, highlight and go to next screen.
        %if exist('rating','var')%BCOE
        DrawFormattedText(window,sprintf('%d',rating),(rect_box(1,rating)+rect_box(2,rating))/2,bounds(4)+small_text,[255, 255, 255]);
        %Screen(window,'Flip'); pause(0.35);
        % break;
        %end
        Screen(window,'Flip');
        pause(0.35);
        
        %% give feedback (correct/incorrect guess)
        % clear screen; draw response
        Screen(window,'TextSize',300);
        if reward
            DrawFormattedText(window, '+','center', 'center', [255,255,255],[],[],[],2);
        else
            DrawFormattedText(window, 'X', 'center', 'center',[255,255,255],[],[],[],2);
        end
        Screen(window,'Flip');
        Screen(window,'TextSize',small_text);
        
        % record data
        color_str='';
        for i=1:colors_shown; color_str=sprintf('%s%c',color_str,c_n(colors(i))); end
        trials{str}=sprintf('%d\t%d\t%d\t%d\t%s',trial,Rexp,reward,rating,color_str);
        if print_results
            fprintf('%d\t\t%d\t\t%d\t\t%d\t\t%s\n',trial,Rexp,reward,rating,color_str);
        end
        ratings(trial,phase)=rating; clear rating;
        str=str+1;
        
        % wait for mouse click to continue on
        [x,y,clicks]=GetMouse(window);
        while ~clicks; [x,y,clicks]=GetMouse(window); end
        pause(0.35);
    end
    
    % setting up prediction phase box sizes
    dL2=(0.75*(rect(3)-rect(1))-n_colors*20)/(n_colors);
    rect_box2=[.12*(rect(3)-rect(1))+(0:n_colors-1)*(dL2+20);ones(1,n_colors)*0.5*(rect(4)-dL2);0.12*(rect(3)-rect(1))+(1:n_colors)*(dL2)+(0:n_colors-1)*20;ones(1,n_colors)*0.5*(rect(4)+dL2)];
    pause(0.35);
    
    %% post phase prediction period
    while 1 % until DONE is clicked
        % display screen elements
        [~,~,bounds]=DrawFormattedText(window, 'Click on the color or colors you think lead \nto receiving a reward. When finished, click DONE', 'center', 0.3*(rect(4)-dL2),[255,255,255],[],[],[],2);
        Screen('FillRect',window,col_spec(:,1:n_colors),rect_box2);
        if any(guesses(phase,:))
            Screen('FrameRect',window,zeros(3,length(c)),[rect_box2(1,c)-3; rect_box2(2,c)-3; rect_box2(3:4,c)],5);
        end
        [~,~,done_box]=DrawFormattedText(window,'DONE','center',rect_box2(1,4)+3*small_text,[255,255,255]);
        [~,~,erase_box]=DrawFormattedText(window,'ERASE','center',rect_box2(1,4)+6*small_text,[255,255,255]);
        Screen(window,'Flip');
        
        % wait for user to select one (or more) colors for their prediction
        % or DONE, record/process response
        pause(0.25);
        [x,y,clicks]=GetMouse(window);
        while ~clicks; [x,y,clicks]=GetMouse(window); end
        
        % if a color is selected
        if y>=0.5*((rect(4)-rect(2))-dL2) && y<=0.5*((rect(4)-rect(2))+dL2)
            for i=1:n_colors
                if x>=rect_box2(1,i) && x<=rect_box2(3,i)
                    guesses(phase,i)=~guesses(phase,i); % toggle guess (select/deselect)
                    [~,c]=find(guesses(phase,:)==1); % record guess
                    break;
                end
            end
            % if done is selected
        elseif x>=done_box(1)&& x<=done_box(3) && y>=done_box(2) && y<=done_box(4)
            break;
            % if erase is selected, clear guesses
        elseif x>=erase_box(1)&& x<=erase_box(3) && y>=erase_box(2) && y<=erase_box(4);
            guesses(phase,:)=zeros(1,n_colors); % start prediction over (solely for ease of use
        end
    end
    pause(0.25);
    
    % record results
    color_str='';
    if exist('c','var')
        for i=1:length(c); color_str=sprintf('%s%c',color_str,c_names(c(i))); end
    end
    trials{str}=sprintf('PREDICTED COLORS: %s', color_str); str=str+1;
end
cols=cols(:,1:n_colors);

%% close screen
Screen('CloseAll');

% meta_data_form is a GUI and records meta data about the test subject.
% Information is temporarily stored in the file meta_data.mat
meta_data_form;
uiwait(findall(0,'type','figure'));
load meta_data.mat;
switch age
    case 1; age_str='16-19';
    case 2; age_str='20-25';
    case 3; age_str='26-35';
    otherwise; age_str='35-50';
end


if print_results
    disp('=====================================================');
end
%% saving data
% saves game data to .mat file and experiment output to .txt file
if save_results
    fprintf('This experiment output has been saved in the file %s.txt',filename');
    save(sprintf('%s',filename),'ratings','guesses','c_n','cols','r_exp_record','r_giv_record');
    txtfile=fopen(sprintf('%s.txt',filename),'wt');
    fprintf(txtfile,'\t\t\t\tCOLOR KEY\n');
    fprintf(txtfile,'tY yellow       G green        B blue\n');
    fprintf(txtfile,'tV violet       P pink         R red\n');
    fprintf(txtfile,'N brown        W white        K black\n\n');
    fprintf(txtfile,'================== EXPERIMENT RESULTS ====================\n\n');
    fprintf(txtfile,'Stimulus A: %c\nStimulus B: %c\n',c_n(1),c_n(2));
    fprintf(txtfile,'Stimulus C: %c\nStimulus D: %c\n',c_n(3),c_n(4));
    fprintf(txtfile,'Stimulus E: %c\nStimulus F: %c\n',c_n(5),c_n(6));
    fprintf(txtfile,'Stimulus G: %c\nStimulus H: %c\n',c_n(7),c_n(8));
    fprintf(txtfile,'Stimulus X: %c\n\n',c_n(9));
    fprintf(txtfile,'==========================================================\n');
    for i=1:str-1; fprintf(txtfile,'%s\n',trials{i}); end
    fprintf(txtfile,'\n\n====================== USER METADATA =====================\n\n');
    fprintf(txtfile,'Identifier:\t%s\n',filename);
    fprintf(txtfile,'Gender:\t\t%s\n',gender);
    fprintf(txtfile,'Age:\t\t%s\n',age_str);
    fprintf(txtfile,'Education:\t%s\n',edu);
    fprintf(txtfile,'Handedness:\t%s\n',handedness);
    fprintf(txtfile,'Feedback:\n');
    if(any(feedback));
        fprintf(txtfile,'%s\n',feedback);
    else
        fprintf(txtfile,'---\n');
    end
    fprintf(txtfile,'==========================================================\n');
end
